import numpy as np
import cv2

COMM = None
try:
  import serial
  COMM = serial.Serial('/dev/ttyAMA0', 115200, timeout=0.5)
except ImportError:
  print('error occured when import serial')

cap = cv2.VideoCapture(0)
# cap.set(3,320)
# cap.set(4,240)
# cap.set(3, 1920.0)
# cap.set(4, 1440.0)
cap.set(3,640)
cap.set(4,480)

fgbg = cv2.createBackgroundSubtractorMOG2(detectShadows=False)

def start():
  while(1):
    ret, frame = cap.read()
    if not ret:
      break
    fgmask = fgbg.apply(frame)
    res = cv2.bitwise_and(frame, frame, mask=fgmask)
    # res = cv2.bilateralFilter(res,9,75,75)
    # fire_frame = check_color(res)
    fire_frame = check_roi(res)
    draw_rectangle(frame, fire_frame)
    cv2.imshow('frame',frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

  cap.release()
  cv2.destroyAllWindows()

def check_roi(img):
  e1 = cv2.getTickCount()
  gray = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
  # gray = cv2.medianBlur(gray, 5)
  ret, binary = cv2.threshold(gray,127,255,cv2.THRESH_BINARY)
  image, contours, hierarchy = cv2.findContours(binary,cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE )
  # if len(contours) > 10:
  #   return img
  for c in contours:
    (x, y, w, h) = cv2.boundingRect(c)
    if c.size > 300:
      roi = img[y:y+h, x:x+w]
      fire_roi = check_color(roi)
      img[y:y+h, x:x+w] = fire_roi
    else:
      img[y:y+h, x:x+w] = 0
  e2 = cv2.getTickCount()
  time = (e2 - e1)/ cv2.getTickFrequency()
  print('check_roi elapsed：', time)
  return img

def check_color(img):
  fire_img = np.zeros(img.shape, np.uint8)
  red_threshold = 55
  saturation_threshold = 115
  b, g, r = cv2.split(img)
  rows, cols, n = img.shape
  for i in range(rows - 1):
    for j in range(cols-1):
      # B,G,R = img[i, j]
      B = img.item(i,j,0)
      G = img.item(i,j,1)
      R = img.item(i,j,2)
      max_value = max((B,G,R))
      min_value = min((B,G,R))
      sum_rgb = sum((B,G,R))
      if sum_rgb == 0:
        S = 0
      else:
        S = (1 - 3*min_value / sum_rgb)
      if R > red_threshold and R >= G and G > B and S >((255-R)*saturation_threshold/red_threshold):
        # fire_img[i,j]= [255,255,255]
        fire_img.itemset((i,j,0),255)
        fire_img.itemset((i,j,1),255)
        fire_img.itemset((i,j,2),255)
      else:
        # fire_img[i,j] = [0, 0, 0]
        fire_img.itemset((i,j,0),0)
        fire_img.itemset((i,j,1),0)
        fire_img.itemset((i,j,2),0)
  return fire_img

def draw_rectangle(img, fire_img):
  e1 = cv2.getTickCount()
  gray = cv2.cvtColor(fire_img,cv2.COLOR_BGR2GRAY)
  # gray = cv2.medianBlur(gray, 5)
  ret, binary = cv2.threshold(gray,127,255,cv2.THRESH_BINARY)
  image, contours, hierarchy = cv2.findContours(binary,cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)

  if len(contours) > 0:
      write_to_serial('1')
  for c in contours:
    # if cv2.contourArea(c) < 1:
    #   continue
    (x, y, w, h) = cv2.boundingRect(c)
    cv2.rectangle(img, (x-5, y-5), (x + 5 + w, y+5 + h), (0, 0, 255), 2)
  # cv2.drawContours(img,contours,-1,(0,255,0),2) 
  e2 = cv2.getTickCount()
  time = (e2 - e1)/ cv2.getTickFrequency()
  print('draw_rectangle elapsed：', time)

def write_to_serial(flag):
  if COMM is not None and COMM.isOpen():
    COMM.write(flag.encode())

if __name__ == '__main__':
  start()